#include "PSDecoder.h"

#include "../inc/Utils.h"

#ifndef AV_RB16
#   define AV_RB16(x)                           \
	((((const unsigned char*)(x))[0] << 8) |          \
	((const unsigned char*)(x))[1])
#endif


typedef struct optional_pes_header{
	unsigned char original_or_copy:1;
	unsigned char copyright:1;
	unsigned char data_alignment_indicator:1;
	unsigned char PES_priority:1;
	unsigned char PES_scrambling_control:2;
	unsigned char fix_bit:2;

	unsigned char PES_extension_flag:1;
	unsigned char PES_CRC_flag:1;
	unsigned char additional_copy_info_flag:1;
	unsigned char DSM_trick_mode_flag:1;
	unsigned char ES_rate_flag:1;
	unsigned char ESCR_flag:1;
	unsigned char PTS_DTS_flags:2;

	unsigned char PES_header_data_length;
}optional_pes_header_t;
static inline uint64_t ff_parse_pes_pts(const unsigned char* buf) {
	return (uint64_t)(*buf & 0x0e) << 29 |
		(AV_RB16(buf+1) >> 1) << 15 |
		AV_RB16(buf+3) >> 1;
}

static uint64_t get_pts(optional_pes_header* option)
{
	if(option->PTS_DTS_flags != 2 && option->PTS_DTS_flags != 3 && option->PTS_DTS_flags != 0)
	{
		return 0;
	}
	if((option->PTS_DTS_flags & 2) == 2)
	{
		unsigned char* pts = (unsigned char*)option + sizeof(optional_pes_header);
		return ff_parse_pes_pts(pts);
	}
	return 0;
}

PSDecoder::PSDecoder(void)
{
	PSBuffer = (unsigned char*)malloc(MAX_FRAME_SIZE);
	memset(PSBuffer, 0, MAX_FRAME_SIZE);
	//PSBuffer = (unsigned char*)malloc(MAX_FRAME_SIZE);
	//memset(PSBuffer, 0, MAX_FRAME_SIZE);
	RTPHeaderSize = 12;
	PSBufferOffset = 0;
	PESFrameOffset = 0;
	bufferSize = 0;
	timestamp = 0;

	recvThreadRunning = true;
	decodeThreadRunning = true;
	playFlag = PLAY;

	recvBuf = (char*)malloc(2048);

	av_register_all();
	avformat_network_init();

	lfds611_queue_new( &qs_pool, QUEUE_SIZE );
	lfds611_queue_new( &qs_packet, QUEUE_SIZE );
	for(int i = 0;i<QUEUE_SIZE;i++)
	{
		PSPacketBuffer *packetBuffer =(PSPacketBuffer*)malloc(sizeof(PSPacketBuffer));;
		packetBuffer->packetBuffer =(unsigned char*)malloc(MAX_PAKCET_SIZE);
		memset(packetBuffer->packetBuffer, 0 , MAX_PAKCET_SIZE);
		lfds611_queue_use(qs_pool);
		lfds611_queue_enqueue( qs_pool, (void *)packetBuffer );
	}

	currentPESFrame = (unsigned char*)malloc(MAX_FRAME_SIZE);
	memset(currentPESFrame, 0, MAX_FRAME_SIZE);
	PESBufferFrame = (unsigned char*)malloc(MAX_FRAME_SIZE);
	memset(PESBufferFrame, 0, MAX_FRAME_SIZE);

	PESFrame = (unsigned char*)malloc(MAX_FRAME_SIZE);
	memset(PESFrame, 0, MAX_FRAME_SIZE);

	PESsplitter = (unsigned char*)malloc(4);
	memset(PESsplitter, 0 , 4);
	PESsplitter[2] = 0x01;
	PESsplitter[3] = 0xe0;

	H264splitter = (unsigned char*)malloc(4);
	memset(H264splitter, 0 , 4);
	H264splitter[3] = 0x01;
}


PSDecoder::~PSDecoder(void)
{
	if(port != 0)
	{
		_print_log("Finished receiving. Closing socket.\n");
		close(RecvSocket);
	}
	if(recvThread != NULL)
	{
		//TerminateThread(startThread, 0);
		recvThreadRunning = false;
		pthread_join(recvThread, NULL);
		recvThread = NULL;
	}

	if(decodeThread != NULL)
	{
		//TerminateThread(decodeThread, 0);
		decodeThreadRunning = false;
		pthread_join(decodeThread, NULL);
		decodeThread = NULL;
	}

	free(currentPESFrame);
	free(PESFrame);
	free(recvBuf);

	if(c != NULL)
	{
		avcodec_close(c);
		av_free(c);
	}

	av_frame_free(&frame);

	if(frameBGR != NULL)
		avpicture_free((AVPicture*) frameBGR);
	avformat_network_deinit();

	if(qs_pool != NULL && qs_packet != NULL)
	{
		PSPacketBuffer *buffer = NULL;
		int free_count = 0;
		lfds611_queue_use(qs_pool);

		while(LFDS_SUCCESS == lfds611_queue_dequeue( qs_pool, (void **) &buffer))
		{
			if(buffer->packetBuffer != NULL)
			{
				free(buffer->packetBuffer);
				free_count ++;
				//printf("pool_free_count:%d\n", free_count);
			}
		}
		lfds611_queue_use(qs_packet);
		while(LFDS_SUCCESS == lfds611_queue_dequeue( qs_packet, (void **) &buffer))
		{
			if(buffer->packetBuffer != NULL)
			{
				free(buffer->packetBuffer);
				free_count ++;
				//printf("packet_free_count:%d\n", free_count);
			}
		}
	}
	else
	{// no need for freeing queue;
	}

	_print_log("released\n");
}


void PSDecoder::startRecvThread()
{
	while(recvThreadRunning == true)
	{
		if(playFlag == PLAY)
		{
		}
		else if(playFlag == PAUSE)
		{
			usleep(NORMAL_WAIT_TIME);
			continue;
		}
		else if(playFlag == STOP)
		{
			break;
		}
		else
		{
			_print_log("illegal playflag\n");
			break;
		}

		int readLength = 0;

		//_print_log("Receiving datagrams...\n");
		struct sockaddr_in clientAddr;
		socklen_t len = sizeof(clientAddr);
		readLength = recvfrom(RecvSocket, recvBuf, 1500, 0, (struct sockaddr*)&clientAddr, &len);

		if (readLength == -1) {
			perror( "recvfrom ");
			//				_print_log("recvfrom failed with error %d\n", WSAGetLastError());
			continue;
		}
		else if(readLength == 0)
		{
			_print_log("End of stream.\n");

			break;
		}
		else
		{

		}

		readLength -= RTPHeaderSize;
		progress += readLength;
		PSPacketBuffer *buffer;
		lfds611_queue_use(qs_pool);
		while(LFDS_SUCCESS != lfds611_queue_dequeue( qs_pool, (void **) &buffer))
		{
			if(recvThreadRunning != true)
				return;
			printf("dequeue failed\n");
			_print_log("dequeue failed\n");
			//lostPacket ++;
			usleep(NORMAL_WAIT_TIME);
		}
		memcpy(buffer->packetBuffer, recvBuf + RTPHeaderSize, readLength);
		buffer->bufferSize = readLength;
		lfds611_queue_use(qs_packet);
		while(LFDS_SUCCESS != lfds611_queue_enqueue( qs_packet, (void *) buffer))
		{
			if(recvThreadRunning != true)
				return;
			lostPacket ++;
			usleep(NORMAL_WAIT_TIME);
		}
	}
}

void PSDecoder::startDecodeThread()
{
	while(qs_pool == NULL || qs_packet == NULL)
		usleep(NORMAL_WAIT_TIME * 100);

	av_init_packet(&avpkt);

	codec = avcodec_find_decoder(AV_CODEC_ID_H264);
	if(codec == NULL)
	{
		_print_log("Codec not found.\n");
	}
	c = avcodec_alloc_context3(codec);
	if(c == NULL)
	{
		_print_log("could not allocate video codec context\n");
	}
	if(avcodec_open2(c, codec, NULL) < 0)
	{
		_print_log("could not open codec\n");
	}
	frame = av_frame_alloc();
	if(frame == NULL)
	{
		_print_log("could not allocate video frame\n");
	}
	frameBGR = av_frame_alloc();
	if(frameBGR == NULL)
	{
		_print_log("could not allocate video frame\n");
	}


	int currentPacketSize = 0;
	char* tempBuffer = NULL;
	while(decodeThreadRunning == true)
	{
		if(playFlag == PLAY)
		{
		}
		else if(playFlag == PAUSE)
		{
			usleep(NORMAL_WAIT_TIME);
			continue;
		}
		else if(playFlag == STOP)
		{
			break;
		}
		else
		{
			_print_log("illegal playflag\n");
			break;
		}

		ESPacketBuffer *buffer;
		lfds611_queue_use(qs_packet);
		while(LFDS_SUCCESS != lfds611_queue_dequeue( qs_packet, (void **) &buffer))
		{
			if(decodeThreadRunning != true)
				return;
			usleep(NORMAL_WAIT_TIME);
		}
		//printf("PSFrameOffset:%d, bufferSize:%d\n", PSFrameOffset,buffer->bufferSize );
		if(PSBufferOffset > 900 * 1024)
		{// packet lost too much, barely to demux, reset the frameoffset to prevent the memory error.
			PSBufferOffset = 0;
			_print_log("packet lost too much, reset the buffer.\n");
		}
		memcpy(PSBuffer + PSBufferOffset, buffer->packetBuffer, buffer->bufferSize);
		//printf("%d\n", buffer->bufferSize);
		PSBufferOffset += buffer->bufferSize;
		currentPacketSize = buffer->bufferSize;
		//printf("offset:%d\n", PSFrameOffset);
		lfds611_queue_use(qs_pool);
		while(LFDS_SUCCESS != lfds611_queue_enqueue( qs_pool, (void *) buffer))
		{
			if(decodeThreadRunning != true)
				return;
			lostPacket ++;
			usleep(NORMAL_WAIT_TIME);
		}
		bool foundPES = false;

		int bodyLength = 0;
		foundPES = findVideoPES(currentPESFrame, PSBuffer, PSBufferOffset, PESsplitter, 4, bodyLength);

		if(foundPES == true)
		{
			if(find(currentPESFrame, bodyLength, H264splitter, 4) == 0)
			{// packet with h264 start code, previous frame ended, decode it.
				if(PESFrameOffset == 0)
				{// previous frame is not the large frame
					memcpy(PESFrame, PESBufferFrame, bufferSize);
					avpkt.data = PESFrame;
					avpkt.size = bufferSize;
					memcpy(PESBufferFrame, currentPESFrame, bodyLength);
					bufferSize = bodyLength;
				}
				else
				{// previous frame is the large frame
					avpkt.data = PESFrame;
					avpkt.size = PESFrameOffset;
					PESFrameOffset = 0;
					memcpy(PESBufferFrame, currentPESFrame, bodyLength);
					bufferSize = bodyLength;
				}
			}
			else
			{// packet without the h264 start code, attach it to the last frame.
				if(PESFrameOffset == 0)
				{
					memcpy(PESFrame + PESFrameOffset, PESBufferFrame, bufferSize);
					PESFrameOffset += bufferSize;
				}

				memcpy(PESFrame + PESFrameOffset, currentPESFrame, bodyLength);
				PESFrameOffset += bodyLength;
				continue;
			}
		}
		else
		{// no complete packet found
			continue;
		}

		int len = 0, got_frame = 0;

		len = avcodec_decode_video2(c, frame, &got_frame, &avpkt);
		if(len < 0)
		{
			_print_log("Error while decoding frame\n");
			continue;
		}

		if(got_frame)
		{
			//_print_log("%d, %d\n", width, height);
			//_print_log("success\n");
			//FILE* fd = fopen("d://1.pgm", "w");
			//fprintf(fd, "P5\n%d %d\n%d\n", width,height,255);
			//for(int i =0;i<height;i++)
			//	fwrite(frame->data[0] + i* frame->linesize[0] , 1, width, fd);
			//fclose(fd);
			if(frameBGR->data[0] == NULL)
			{
				width = frame->width;
				height = frame->height;
				pRGBData = (char*)malloc(width * height * 4);
				pSwsCtxBGR32 = sws_getContext(frame->width, frame->height,AV_PIX_FMT_YUV420P, width, height, AV_PIX_FMT_RGB32,SWS_BILINEAR,NULL,NULL,NULL);

				avpicture_alloc((AVPicture*)frameBGR, AV_PIX_FMT_RGB32, width, height);
			}
			frameBGR->width = width;
			frameBGR->height = height;
			int ret = sws_scale(pSwsCtxBGR32, frame->data,frame->linesize,0, frame->height,frameBGR->data,frameBGR->linesize);
			if(frameBGR->height != ret)
			{
				_print_log("sws_scale failed\n");
				continue;
			}
			frameInfo.timeStamp = timestamp;
			frameInfo.width = width;
			frameInfo.height = height;
			(*DecodeCBFun)((char*)frameBGR->data[0], width * height * 4, &frameInfo, pUser);

		}
		av_free_packet(&avpkt);		
	}


}

int PSDecoder::find(unsigned char* buffer,  int bufferLength, unsigned char* splitter,int splitterLength)
{
	int retval = -1;
	int pos = 0;
	int index = 0;
	while(pos != bufferLength)
	{
		for(int i = 0; i<splitterLength;i++)
		{
			if(buffer[pos + i] == splitter[i])
				index++;
			else
				break;
		}
		if(index == splitterLength)
		{
			retval = pos;
			break;
		}

		else
		{
			index = 0;
			pos++;
		}
	}
	return retval;
}

/* PES packet header
0               1               2               3
0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7 0
+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+
|				  packet start code	              |
+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+
|   stream id   |      PES Packet length		  |
+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+
Optional PES header                
|   |   | | | | |   | | | | | | | header length |
+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+
*/



bool PSDecoder::findVideoPES(unsigned char* currentPESFrame, unsigned char* PSBuffer,  int& bufferLength, unsigned char* splitter,int splitterLength, int& bodyLength)
{
	bool retval = false;
	int PESStart = find(PSBuffer,  bufferLength, splitter, splitterLength);
	if(PESStart < 0)
		retval = false;
	else
	{
		if(PESStart + 8 > bufferLength)
			retval = false;
		else
		{
			int optionalHeaderLength = 3;
			int PESPacketHeaderLength = 6;

			optional_pes_header* optionalHeader = (optional_pes_header*)(PSBuffer + PESStart + PESPacketHeaderLength);
			timestamp = get_pts(optionalHeader) / 90;

			unsigned char packetLengthHigh = PSBuffer[PESStart + 4];
			unsigned char packetLengthLow = PSBuffer[PESStart + 5];
			int packetLength = (packetLengthHigh << 8) | packetLengthLow; //from byte 6


			int PESHeaderLength = PSBuffer[PESStart + PESPacketHeaderLength + 2]; 

			bodyLength = (packetLength - PESHeaderLength) - 3;
			//printf("packetLength:%d, bodyLength:%d\n", packetLength, bodyLength);
			int bodyOffset = PESStart + PESHeaderLength + PESPacketHeaderLength + optionalHeaderLength;
			if (bufferLength < (packetLength + PESPacketHeaderLength))
			{
				retval = false;
			}
			else
			{
				int pos_end_ret = (PESStart + packetLength) + 6;
				if (pos_end_ret >= bufferLength)
				{
					retval = false;
				}
				else
				{
					if (bodyLength <= 0)
					{
						memcpy(PSBuffer + pos_end_ret, PSBuffer, bufferLength - pos_end_ret);
						bufferLength -= pos_end_ret;
						retval = false;
					}
					else
					{
						memcpy(currentPESFrame, PSBuffer + bodyOffset, bodyLength);
						memcpy(PSBuffer, PSBuffer + pos_end_ret, bufferLength - pos_end_ret);
						bufferLength -= pos_end_ret;
						retval = true;
					}
				}
			}
		}
	}


	return retval;
}
